import torch
import torch.nn as nn

from stft import STFTEncoder, STFTDecoder
from net import IntraFrame, InterFrame, MultiHeadAttention


class TF_GridNet(nn.Module):
    def __init__(
        self,
        emb_dim=32,
        kernel_size=3,
        eps=1e-5,
        n_fft=256,
        B=6,
        src=2,
        hidden_channels=192,
        device="cuda",
    ):
        super().__init__()
        self.encoder = STFTEncoder(n_fft=n_fft, device=device)
        self.frequency_dim = n_fft // 2 + 1  # 129
        self.conv2d = nn.Conv2d(
            in_channels=2,
            out_channels=emb_dim,
            kernel_size=(kernel_size, kernel_size),
            padding=(kernel_size // 2, 1),
        )
        self.group_norm = nn.GroupNorm(1, emb_dim, eps=eps)

        self.blocks = []
        for _ in range(0, B):
            self.blocks.append(
                IntraFrame(emb_channel=emb_dim, hidden_channels=hidden_channels)
            )
            self.blocks.append(
                InterFrame(emb_channel=emb_dim, hidden_channels=hidden_channels)
            )
            self.blocks.append(MultiHeadAttention(head=1, emb_dim=emb_dim))
        self.tf_grid_block = nn.Sequential(*self.blocks)
        self.dconv = nn.ConvTranspose2d(emb_dim, 2 * src, (3, 3), padding=(3 // 2, 1))
        self.src = src

        self.decoder = STFTDecoder(device=device)

    def forward(self, input):
        # [B, 2, T, F]
        x = self.encoder(input)

        # [B, emb, T, F]
        x = self.group_norm(self.conv2d(x))

        # main TF_grid block
        x = self.tf_grid_block(x)

        # deconv [B, 2 * src, T, F]
        x = self.dconv(x)

        B, _, T, F = x.shape
        x = x.view(B, self.src, 2, T, F)
        x = torch.complex(x[:, :, 0], x[:, :, 1])
        x = x.view(B * self.src, T, F)

        # stft inverse
        x = self.decoder(x, B)

        return x


if __name__ == "__main__":
    tf_gridnet = TF_GridNet()
    pytorch_total_params = sum(p.numel() for p in tf_gridnet.parameters())
    print(pytorch_total_params)  # 8.1M
    x = torch.randn(2, 1, 32000)
    y = tf_gridnet(x)
    print(y.shape)
    pass
